#!/usr/bin/env python3
"""
Launch script for OmniDet MTL.

# usage: ./main.py --config data/params.yaml

# author: Varun Ravi Kumar <rvarun7777@gmail.com>

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; Authors provide no warranty with the software
and are not liable for anything.
"""

import argparse
import json
import os
import shutil

from debian.debtags import output
# from distutils.util import strtobool
from setuptools._distutils.util import strtobool
from pathlib import Path
import logging

import yaml
# from ruamel.yaml import YAML
from omnidet.utils import Tupperware


def printj(dic):
    return print(json.dumps(dic, indent=4))


def collect_args() -> argparse.Namespace:
    """Set command line arguments"""
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', help="Config file", type=str, default=Path(__file__).parent / "data/params.yaml")
    args = parser.parse_args()
    return args


def collect_tupperware() -> Tupperware:
    config = collect_args()
    params = yaml.safe_load(open(config.config))
    args = Tupperware(params)
    printj(args)
    return args


def main():
    args = collect_tupperware()
    # 日志的输出等级确认
    logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s')
    # 当前工程路径
    working_path = os.path.dirname(__file__)
    # 设置预训练模型的路径
    pretrained_dir_path = os.path.join(working_path, 'pre-trained/weights')
    if os.path.exists(pretrained_dir_path) is False:
        os.makedirs(pretrained_dir_path)
        logging.error('pretrained_path is not exist, create it! please check the pretrained weights now!')
        assert False
    pretrained_weights_path = os.path.join(pretrained_dir_path, 'res' + str(args.network_layers))
    args.pretrained_weights = pretrained_weights_path
    # 设置dataset的路径
    dataset_dir = os.path.join(working_path, args.dataset_dir)
    args.dataset = dataset_dir
    args.train_file = os.path.join(working_path, args.train_file)
    args.val_file = os.path.join(working_path, args.val_file)
    args.test_file = os.path.join(working_path, args.test_file)
    # 设置训练模型的输出路径
    output_dir_path = os.path.join(working_path, args.output_directory)
    if os.path.exists(output_dir_path) is False:
        os.makedirs(output_dir_path)
        logging.error('output_dir_path is not exist, create it! ')
        assert False
    args.output_directory = output_dir_path
    # 设置标签路径及文件名


    # log_path = os.path.join(args.output_directory, args.model_name)

    # if os.path.isdir(log_path):
    #     # pass
    #     if strtobool(input("=> Clean up the log directory?")):
    #         shutil.rmtree(log_path, ignore_errors=False, onerror=None)
    #         os.mkdir(log_path)
    #         print("=> Cleaned up the logs!")
    #     else:
    #         print("=> No clean up performed!")
    # else:
    #     print(f"=> No pre-existing directories found for this experiment. \n"
    #           f"=> Creating a new one!")
    #     os.mkdir(log_path)

    # os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    # os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices or "-1"

    if args.train == "distance":
        model = DistanceModel(args)
        model.distance_train()
    elif args.train == "semantic":
        model = SemanticModel(args)
        model.semantic_train()
    elif args.train == "motion":
        model = MotionModel(args)
        model.motion_train()
    elif args.train == "detection":
        model = DetectionModel(args)
        model.detection_train()
        # model.detection_train()
    elif args.train == "distance_semantic":
        model = DistanceSemanticModel(args)
        model.distance_semantic_train()
    elif args.train == "detection_semantic":
        model = DetectionSemanticModel(args)
        model.detection_semantic_train()
    elif args.train == "distance_semantic_motion":
        model = DistanceSemanticMotionModel(args)
        model.distance_semantic_motion_train()
    elif args.train == "distance_semantic_detection":
        model = DistanceSemanticDetectionModel(args)
        model.distance_semantic_detection_train()
    elif args.train == "distance_semantic_detection_motion":
        model = DistanceSemanticDetectionMotionModel(args)
        model.distance_semantic_detection_motion_train()
    else:
        raise NotImplementedError


if __name__ == "__main__":
    from train_detection import DetectionModel
    from train_distance import DistanceModel
    from train_distance_semantic import DistanceSemanticModel
    from train_distance_semantic_detection import DistanceSemanticDetectionModel
    from train_distance_semantic_detection_motion import DistanceSemanticDetectionMotionModel
    from train_distance_semantic_motion import DistanceSemanticMotionModel
    from train_detection_semantic import DetectionSemanticModel
    from train_motion import MotionModel
    from train_semantic import SemanticModel

    main()
